import os
import json
import cv2
from PIL import Image
from ultralytics import YOLO
import extract_seal as es
from collections import defaultdict


class YOLOPredictionProcessor:
    def __init__(self, model_path, extract_folder="cropped_seals"):
        """
        初始化YOLO预测处理器

        参数:
            model_path: 训练好的YOLO模型路径(.pt文件)
            extract_folder: 裁剪印章保存目录
        """
        self.model_path = model_path
        self.extract_folder = extract_folder
        self.model = None
        self.class_mapping = {
            0: "red_off_seal",
            1: "red_person_seal",
            2: "black_off_seal",
            3: "black_person_seal"
        }

        # 初始化模型和输出目录
        self._init_model()
        os.makedirs(self.extract_folder, exist_ok=True)

    def _init_model(self):
        """加载YOLO模型"""
        print(f"加载YOLO模型: {self.model_path}")
        self.model = YOLO(self.model_path)

    def process_predictions(self, input_folder, save_dict_json=False)->dict:
        """
        处理YOLO预测结果

        参数:
            input_folder: 包含输入图片的文件夹路径
            save_dict_json: 是否将类别字典保存为JSON文件

        返回:
            dict: 检测结果字典 {图片名: [检测到的类别ID列表]}
        """
        print(f"开始处理YOLO预测，输入目录: {input_folder}")
        result_dict = {}

        for img_name in self._get_image_files(input_folder):
            img_path = os.path.join(input_folder, img_name)
            base_name = os.path.splitext(img_name)[0]

            # 执行预测并处理结果
            detected_classes = self._process_single_image(img_path, base_name)
            result_dict[img_name] = list(detected_classes)

        # 可选保存结果
        if save_dict_json:
            self._save_results_to_json(result_dict)

        print("YOLO检测处理完成")
        return result_dict

    def _get_image_files(self, input_folder):
        """获取输入目录中的所有图片文件"""
        return [f for f in os.listdir(input_folder)
                if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    def _process_single_image(self, img_path, base_name):
        """处理单张图片的预测结果"""
        results = self.model.predict(img_path)
        detected_classes = set()

        for i, result in enumerate(results):
            boxes = result.boxes
            orig_img = result.orig_img

            for j, box in enumerate(boxes):
                cls_id = int(box.cls)
                detected_classes.add(cls_id)

                # 处理每个检测框
                self._process_detection_box(
                    orig_img,
                    box.xyxy[0],
                    base_name,
                    cls_id,
                    j
                )

        return detected_classes

    def _process_detection_box(self, orig_img, box_coords, base_name, cls_id, box_idx):
        """处理单个检测框"""
        x1, y1, x2, y2 = map(int, box_coords)
        cropped_img = orig_img[y1:y2, x1:x2]

        # if cropped_img.shape[-1] == 3:  # 如果是彩色图像
        #     cropped_img = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)

        # 印章提取处理
        cropped_img = self._extract_seal(cropped_img)

        if cropped_img.shape[-1] == 3:  # 如果是彩色图像
            cropped_img = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)

        # 生成保存路径
        save_name = f"{base_name}-{self.class_mapping[cls_id]}-{cls_id}-{box_idx}.jpg"
        save_path = os.path.join(self.extract_folder, save_name)

        # 保存图像
        self._save_cropped_image(cropped_img, save_path)

    def _extract_seal(self, image):
        """印章提取处理"""
        # 这里调用你的es.extract_seal_array方法
        return es.extract_seal_array(image, h_scale=[150, 30], s_scale=[20, 255], v_scale=[20, 255])


        # 如果没有es模块，可以保留原始图像或实现自己的提取逻辑
        # if image.shape[-1] == 3:  # 如果是彩色图像
        #     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # return image

    def _save_cropped_image(self, image_array, save_path):
        """保存裁剪后的图像"""
        Image.fromarray(image_array).save(save_path)

    def _save_results_to_json(self, result_dict, filename="detection_results.json"):
        """保存结果到JSON文件"""
        with open(filename, "w") as f:
            json.dump(result_dict, f, indent=4)
        print(f"检测结果已保存到: {filename}")


if __name__ == "__main__":
    # 使用示例
    processor = YOLOPredictionProcessor(
        model_path="seal_sign_model.pt",
        extract_folder="cropped_seals"
    )

    results = processor.process_predictions(
        input_folder="seal_examples",
        save_dict_json=True
    )
